import datetime
import os

import torch
import whisper
from loguru import logger

# 加载模型，'base'是较小较快的模型，还有small, medium, large等更大更准确的模型
model = whisper.load_model("medium", device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))


def generate_srt_file(video_origin_path) -> str:
    """
    根据视频地址生成字幕文件

    :param video_origin_path: 视频地址
    :return: 字幕文件地址
    """
    logger.info(f"开始生成{video_origin_path}字幕文件...")
    if not os.path.exists(video_origin_path):
        raise FileNotFoundError(f"视频文件不存在：{video_origin_path}")

    # 字幕文件地址
    srt_file_path = video_origin_path.replace("videos", "srt").replace(".mp4", ".srt")

    # 识别视频文件中的音频并生成字幕
    result = model.transcribe(video_origin_path)

    # 保存为SRT字幕文件
    with open(srt_file_path, "w", encoding="utf-8") as f:
        for i, segment in enumerate(result["segments"]):
            start = segment["start"]
            end = segment["end"]
            text = segment["text"]
            # 写入SRT格式：序号、时间轴、字幕文本
            f.write(f"{i+1}\n")
            f.write(f"{seconds_to_srt_timestamp(start)} --> {seconds_to_srt_timestamp(end)}\n")
            f.write(f"{text}\n\n")
    logger.info(f"字幕文件已保存：{srt_file_path}")
    return srt_file_path


def seconds_to_srt_timestamp(seconds):
    """
    将秒数转换为SRT格式的时间戳 (HH:MM:SS,mmm)

    参数:
    seconds -- 秒数，可以是浮点数

    返回:
    str -- 格式化后的时间字符串
    """
    # 创建timedelta对象并转换为标准时间格式
    td = datetime.timedelta(seconds=seconds)
    total_seconds = int(td.total_seconds())
    hours, remainder = divmod(total_seconds, 3600)
    minutes, seconds = divmod(remainder, 60)
    milliseconds = int((td.total_seconds() - total_seconds) * 1000)

    # 格式化为 HH:MM:SS,mmm
    return f"{hours:02d}:{minutes:02d}:{seconds:02d},{milliseconds:03d}"

if __name__ == '__main__':
    print(torch.cuda.is_available())